﻿using MultiAgentsServer.Scripts.Shared;
using MultiAgentsShared;
using OpenAI.ObjectModels;
using SharpToken;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using System.Text.RegularExpressions;

namespace MultiAgentsServer.Scripts.Tools;

public class MessageStringTool()
{
    public static string GetAgentMentionFromMessage(string message)
    {
        Regex mentionAgentRegex = new(@"@(\S+)");

        Match mentionMatch = mentionAgentRegex.Match(message);

        string agentTemplateName = mentionMatch.Success ? mentionMatch.Groups[1].Value : NameConsts.DefaultAgentName;

        return agentTemplateName;
    }

    public static int CalculateTokens(string model, string? message = null, Dictionary<string, string>? images = null, Dictionary<string, string>? textFiles = null)
    {
        GptEncoding gptEncoding;

        if (model == Models.EnumToString(Models.Model.Gpt_4o_mini_2024_07_18) || model == Models.EnumToString(Models.Model.Gpt_4o))
        {
            gptEncoding = GptEncoding.GetEncodingForModel("gpt-4o");
        }
        else if (model == Models.EnumToString(Models.Model.Gpt_4_turbo_2024_04_09))
        {
            gptEncoding = GptEncoding.GetEncodingForModel("gpt-4");
        }
        else
        {
            return 0;
        }

        int tokens = 0;

        if (images != null)
        {
            foreach (KeyValuePair<string, string> image in images)
            {
                byte[] imageBytes = Convert.FromBase64String(image.Value);

                Image<Rgba32>? imageItem = Image.Load<Rgba32>(imageBytes);
                int width = imageItem.Width;
                int height = imageItem.Height;
                int tilesWidth = (int)Math.Ceiling(width / (double)512);
                int tilesHigh = (int)Math.Ceiling(height / (double)512);
                tokens += tilesWidth * tilesHigh * 170 + 85;
            }
        }
        
        if (textFiles != null)
        {
            foreach (KeyValuePair<string, string> textFile in textFiles)
            {
                tokens += gptEncoding.CountTokens(textFile.Value);
            }
        }

        if (message != null)
        {
            tokens += gptEncoding.CountTokens(message);
        }

        return tokens;
    }
}